import os
from call_model import call_embedding_model
from database import init_faiss, get_faiss, get_cursor
from read_words import insert_words
from read_pdf import insert_pdf_infos

query_text = "测试text embedding模型的输出维度"
query_embedding = call_embedding_model(query_text)

if query_embedding is None:
    import sys
    print("embedding模型执行错误")
    sys.exit()

faiss_path = "/workspace/database/index.faiss"
dataset_path = '/workspace/database/dataset.db'

# 初始化faiss和sqlite
if not os.path.exists(faiss_path):
    index = init_faiss(faiss_path, query_embedding.shape[1])
else:
    index = get_faiss(faiss_path)

conn, cursor = get_cursor(dataset_path)


if False:
    filepath = '/workspace/Data/终稿_附件_2_：人文关怀类假期20240715.doc'

    result = insert_words(filepath, conn, cursor, index, faiss_path)

    if result is None:
        print(result)


if True:
    pdf_path = "/workspace/Data/Audit介绍资料.pdf"
    prompt = "提取总结这张图片中的信息"
    insert_pdf_infos(pdf_path, prompt, index, faiss_path, conn, cursor)